#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
主题: 增量式解析大型XML文件
问题: 你想使用尽可能少的内存从一个超大的XML文档中提取数据。
提示 : 
"""

from xml.etree.ElementTree import iterparse, parse
from collections import Counter

filename = 'potholes.xml'

'''按照坑洼报告数量排列邮编号码'''

def recipe1():
    '''按照坑洼报告数量排列邮编号码: 占用大量内存，但性能好'''
    potholes_by_zip = Counter()

    doc = parse(filename)
    for pothole in doc.iterfind('row/row'):
        potholes_by_zip[pothole.findtext('zip')] += 1
    for zipcode, num in potholes_by_zip.most_common():
        print(f"{zipcode, num = }")

def recipe2():
    '''按照坑洼报告数量排列邮编号码: 大大节约了内存资源，但性能差'''
    potholes_by_zip = Counter()
    data = parse_and_remove(filename, 'row/row')
    for pothole in data:
        potholes_by_zip[pothole.findtext('zip')] += 1
    for zipcode, num in potholes_by_zip.most_common():
        print(f"{zipcode, num = }")

def recipe3():
    data = iterparse(filename,('start','end'))
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")
    print(f"{next(data) = }")

def parse_and_remove(filename, path):
    path_parts = path.split('/')
    doc = iterparse(filename, ('start', 'end'))
    # Skip the root element
    next(doc)

    tag_stack = []
    elem_stack = []
    for event, elem in doc:
        if event == 'start':
            tag_stack.append(elem.tag)
            elem_stack.append(elem)
        elif event == 'end':
            if tag_stack == path_parts:
                yield elem
                elem_stack[-2].remove(elem)     # 这个语句才是使得程序占用极少内存的ElementTree的核心特性
            try:
                tag_stack.pop()
                elem_stack.pop()
            except IndexError:
                pass

def main():
    print('recipe1'.center(20, '*'))
    recipe1()
    print('recipe2'.center(20, '*'))
    recipe2()
    print('recipe3'.center(20, '*'))
    recipe3()

if __name__ == '__main__':
    main()            